flax
jaxlib
tensorflow
tensorflow-datasets
clu
ml_collections
numpy